#!/usr/bin/env python
import sys
import serial
import math
import time
import cv2
import numpy as np
import threading
from kyle_robot_toolbox.camera import USBCamera
from kyle_robot_toolbox.robot_arm.arm5dof_uservo import Arm5DoFUServo
from kyle_robot_toolbox.robot_arm.arm5dof_uservo import Arm5DoFApplication
from image2workspace import Image2Workspace
from color_block_detector import ColorBlockDetector


class xArmGrasp(object):
    def __init__(self, port='/dev/ttyUSB0', baud=115200):
        self.__ser = serial.Serial(port, baud, timeout=0.5)
        self.__ser.bytesize = 8
        self.__ser.parity = serial.PARITY_NONE
        self.__ser.stopbits = 1

        if self.__ser.isOpen():
            self.__ser.close()
        try:
            self.__ser.open()
        except serial.SerialException as msg:
            print('Serial open failed! :' + msg)
            sys.exit()
        
        #Add rx data threading 
        rx_th = threading.Thread(target = self.SerialRecvieCallback)
        rx_th.setDaemon(True)
        rx_th.start()
        
        self.cubic_min_size = 150
        self.cubic_max_size = 400
        self.arm_offset_xy = 0.0 # XY平面上的偏移量
        self.arm_offset_z = 0.0 # Z轴上的偏移量
        self.img_offset_x = 100
        self.cubic_count_threshold = 3 #在某个位置检测到多少次就进行抓取
        self.grab_cubic_in_ws = False # - 是否只抓取工作区内的物块,True: 只抓取工作区内的物块,False: 工作区外的也抓
        self.color_names = ['red'] # 要抓取的色块颜色
        self.grasp_stauts = False
        self.place_stauts = False

        ## 创建机械臂
        arm_config_folder = "./config"
        self.arm = Arm5DoFUServo(config_folder=arm_config_folder)
        self.arm.set_tool_pose(pose_name="capture_image")
        self.arm_app = Arm5DoFApplication(self.arm, config_folder=arm_config_folder)

        # 创建摄像头
        self.camera = USBCamera(config_path="config/usb_camera")
        self.camera.init_video_stream()
        self.camera.load_cam_calib_data()

        self.color_block_detector = ColorBlockDetector(self.color_names)
        self.img2ws = Image2Workspace(self.camera)
        self.candi_workpiece_dict = {} 
        
        try:
            while True:
                if self.grasp_stauts is True:
                    img, px, py, color_name = self.get_object_image_postion()
                    print("image : ", px, py)
                    wx, wy = self.img2ws.img2workspace(px, py)
                    if self.grab_cubic_in_ws:
                        if abs(wx) > img2ws.x0 or abs(wy) > img2ws.y0:
                            # 在工作区格子外，不处理
                            print(f"超出工作区范围: wx = {wx} wy = {wy}")
                            continue 
                    # 对位置进行编码
                    code = self.candi_workpiece_encoding(wx, wy)
                    if code not in self.candi_workpiece_dict:
                        self.candi_workpiece_dict[code] = [color_name, 1]
                    else:
                        self.candi_workpiece_dict[code][0] = color_name
                        self.candi_workpiece_dict[code][1] += 1
                    # print(f"纳入统计: wx = {wx} wy = {wy} 编码 = {code}")
                    if(self.candi_workpiece_dict[code][1] >= self.cubic_count_threshold):
                        print(f"物块颜色 {color_name} wx={wx:.1f} wy={wy:.1f} mm")
                        wx2, wy2, wz2 = self.justify_workspace_posi(wx, wy, self.arm_app.object_height)
                        self.arm_app.grab_cubic(wx2, wy2, wz2)
                        # 清空投票字典
                        self.candi_workpiece_dict = {}
                        self.grasp_stauts = False
                if self.place_stauts is True:
                    self.arm_app.place_cubic()
                    self.place_stauts = False
        except KeyboardInterrupt:
            print("按键中断")
        finally:
            # 关闭所有的窗口
            cv2.destroyAllWindows()
            # 释放相机
            self.camera.release()

    def SerialRecvieCallback(self):
        while True:
            rx_len = self.__ser.inWaiting() 
            if rx_len > 0:
                rx_data = self.__ser.read(rx_len)
                rx_data = rx_data.decode()
                print("json data: ", rx_data)
                if rx_data=='抓取':
                    self.grasp_stauts = True
                if rx_data == "放下":
                    self.place_stauts = True
            
       
    def get_object_image_postion(self):
        img_bgr = self.camera.read_color_img()
        img_bgr = self.camera.remove_distortion(img_bgr)
        canvas = np.copy(img_bgr)
        for color_name in self.color_names:
            rect_list, canvas = self.color_block_detector.color_clock_rect(\
				img_bgr, color_name, canvas=canvas)
            for rect in rect_list:
                x, y, w, h = rect
                if w < self.cubic_min_size or w > self.cubic_max_size:
                    print(f"像素宽度 = {w}, 不满足尺寸要求")
                    continue
                if h < self.cubic_min_size or h > self.cubic_max_size:
                    print(f"像素高度 = {h}, 不满足尺寸要求")
                    continue
                if h > w:
                    px = x + w/2
                    py = y + h - 0.8*(w/2) - self.img_offset_x
                else:
                    px = x + w/2
                    py = y + h/2 -0.8*(w/2) - self.img_offset_x
				# 绘制中心点
                cv2.circle(canvas, (int(px), int(py)), 10, (255, 0, 255), thickness=-1)
                return canvas, px, py, color_name
        return canvas, 0, 0, None

    def candi_workpiece_encoding(self, wx, wy, delta=20):
        #'''工件坐标编码, 坐标单位mm'''
        return int(wx/delta) + int(wy/delta)*1000

    def justify_workspace_posi(self, wx, wy, wz):
        #'''调整工作空间坐标'''
        theta = math.atan2(wy, wx)
        r = math.sqrt(wx*wx + wy*wy)
        new_r = r + self.arm_offset_xy
        new_wx = new_r * math.cos(theta)
        new_wy = new_r * math.sin(theta)
        new_wz = wz + self.arm_offset_z
        return new_wx, new_wy, new_wz

       
    def SerialWriteData(self, data):
        self.__ser.write(data)


if __name__ == '__main__':
    try:
        xArm = xArmGrasp()
    except Exception.message:
        print("...........xxx............"+ Exception.message)
        pass
